use axum::{
    extract::{Json, State},
    http::{HeaderMap, StatusCode},
    response::IntoResponse,
    routing::{get, post},
    Router,
};
use axum_extra::TypedHeader;
use axum_extra::headers::{Authorization, authorization::Bearer};
use bytes::Bytes;
use futures::stream::Stream;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tower_http::cors::CorsLayer;
use tracing::{info, error, warn};
use tower_http::trace::TraceLayer;

use crate::{config::Config, token::get_cached_copilot_token};

pub struct AppState {
    pub config: Config,
}

#[derive(Debug, Serialize)]
struct ModelInfo {
    id: String,
    object: String,
    created: i64,
    owned_by: String,
    permission: Vec<String>,
}

#[derive(Debug, Serialize)]
struct ModelsResponse {
    object: String,
    data: Vec<ModelInfo>,
}

#[derive(Debug, Deserialize, Serialize)]
struct ChatMessage {
    role: String,
    content: serde_json::Value,
}

#[derive(Debug, Deserialize, Serialize)]
struct ChatRequest {
    messages: Vec<ChatMessage>,
    #[serde(flatten)]
    extra: serde_json::Value,
}

pub fn create_router(config: Config) -> Router {
    let state = Arc::new(AppState { config });

    Router::new()
        .route("/v1/models", get(list_models))
        .route("/v1/chat/completions", post(chat_completions))
        .layer(TraceLayer::new_for_http())
        .layer(CorsLayer::permissive())
        .with_state(state)
}

async fn validate_auth_token(
    TypedHeader(auth): TypedHeader<Authorization<Bearer>>,
    State(state): State<Arc<AppState>>,
) -> Result<(), (StatusCode, Json<serde_json::Value>)> {
    if !state.config.auth_tokens.contains(&auth.token().to_string()) {
        return Err((
            StatusCode::UNAUTHORIZED,
            Json(serde_json::json!({
                "error": {
                    "code": "unauthorized",
                    "message": "Invalid authorization token",
                    "details": "Invalid authorization token",
                }
            })),
        ));
    }
    Ok(())
}

async fn list_models(
) -> Result<Json<ModelsResponse>, (StatusCode, Json<serde_json::Value>)> {
    Ok(Json(ModelsResponse {
        object: "list".to_string(),
        data: vec![
            ModelInfo {
                id: "o1".to_string(),
                object: "model".to_string(),
                created: 1687882411,
                owned_by: "github-copilot".to_string(),
                permission: vec![],
            },
            ModelInfo {
                id: "o1-preview".to_string(),
                object: "model".to_string(),
                created: 1687882411,
                owned_by: "github-copilot".to_string(),
                permission: vec![],
            },
            ModelInfo {
                id: "o1-mini".to_string(),
                object: "model".to_string(),
                created: 1687882411,
                owned_by: "github-copilot".to_string(),
                permission: vec![],
            },
            ModelInfo {
                id: "gpt-4o".to_string(),
                object: "model".to_string(),
                created: 1687882411,
                owned_by: "github-copilot".to_string(),
                permission: vec![],
            },
            ModelInfo {
                id: "claude-3.5-sonnet".to_string(),
                object: "model".to_string(),
                created: 1687882411,
                owned_by: "github-copilot".to_string(),
                permission: vec![],
            },
        ],
    }))
}

async fn chat_completions(
    auth: TypedHeader<Authorization<Bearer>>,
    State(state): State<Arc<AppState>>,
    Json(request): Json<ChatRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
    info!("收到聊天请求: {:?}", request);

    validate_auth_token(auth, State(state.clone())).await?;

    let token = get_cached_copilot_token(&state.config)
        .await
        .map_err(|e| {
            error!("获取 Copilot token 失败: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({
                    "error": {
                        "message": e.to_string(),
                        "type": "internal_error",
                        "code": "internal_error",
                        "param": null,
                    }
                })),
            )
        })?;

    info!("成功获取 Copilot token");

    let client = reqwest::Client::new();
    let response = client
        .post(&state.config.api_url)
        .header("Authorization", format!("Bearer {}", token.token))
        .header("Content-Type", "application/json")
        .header("Accept", "text/event-stream")
        .header("editor-version", "vscode/1.95.3")
        .json(&request)
        .send()
        .await
        .map_err(|e| {
            error!("发送请求到 Copilot API 失败: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({
                    "error": {
                        "message": e.to_string(),
                        "type": "internal_error",
                        "code": "internal_error",
                        "param": null,
                    }
                })),
            )
        })?;

    if !response.status().is_success() {
        let status = response.status();
        let error_text = response
            .text()
            .await
            .unwrap_or_else(|_| "Unknown error".to_string());
        warn!("Copilot API 返回错误: {} - {}", status, error_text);
        return Err((
            StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_REQUEST),
            Json(serde_json::json!({
                "error": {
                    "message": error_text,
                    "type": if status == 429 { "rate_limit_exceeded" } else { "api_error" },
                    "code": if status == 429 { "rate_limit_exceeded" } else { "api_error" },
                    "param": null,
                }
            })),
        ));
    }

    info!("成功获取 Copilot API 响应");

    let stream = response.bytes_stream();
    let body = axum::body::Body::from_stream(stream);

    let mut headers = HeaderMap::new();
    headers.insert(
        "Content-Type",
        "text/event-stream".parse().expect("valid header value"),
    );

    Ok((headers, body))
} 